#! -*- coding:utf-8 -*-
# !/env/bin/python3

'''Some utils for yolt
'''

from os import listdir
from os import path
from os.path import join, splitext
import numpy as np
import cv2


def getFilePathsInDir(dir_path, suffix, begin_index=0, num=-1):
    '''Get specified numbers of files whose suffix is 'suffix' in directory
            'dir_path' and return a list of file paths
        Parameters:
            dir_path: directory path
            suffix: specified file suffix
            begin_index: the first file's index
            num: number of read file
        Return [list]: a list of file paths
    '''
    assert begin_index >= 0

    filenames = listdir(dir_path)
    file_paths = []
    # for dirpath, dirnames, filenames in walk(dir_path):
    for filename in filenames:
        if splitext(filename)[1] == suffix:
            file_paths.append(join(dir_path, filename))

    if begin_index >= len(file_paths):
        return []

    file_paths.sort()

    if num == -1:
        return file_paths[begin_index:]

    return file_paths[begin_index:begin_index + num]


def getLabelTxtData(label_txt_path, row, col, depth):
    '''Read label txt file's data, and convert to specified format, return ndarray
        Parameters:
            label_txt_path: label txt file path
            row, col, depth: specified (row, col, channels)
        Return [ndarray]: label data
    '''
    data = np.loadtxt(label_txt_path)
    assert row * col * depth == data.shape[0] * data.shape[1]

    data = np.asarray(data, dtype='float')
    return data.reshape((row, col, depth))


def extractBoundingBoxesFromLabel(label_data, image_width, image_height, box_num, iou_thresh=0.24):
    '''Extract bounding boxes from label data and return a list of bounding boxes
        Parameters:
            label_data: a ndarray of label data
        Return [list]: bounding boxe's coordinate
    '''
    [rows, cols, depth] = label_data.shape
    assert depth % box_num == 0
    assert image_width % cols == 0
    assert image_height % rows == 0

    PARAM_NUM = depth // box_num
    CELL_WIDTH = image_width // cols
    CELL_HEIGHT = image_height // rows

    boxes = []
    for row in range(rows):
        for col in range(cols):
            for box in range(box_num):
                iou = label_data[row, col, box *
                                 PARAM_NUM + 4] / (CELL_WIDTH / 2)
                if iou > iou_thresh:
                    cx = col * CELL_WIDTH + \
                         label_data[row, col, box * PARAM_NUM + 0]
                    cy = row * CELL_HEIGHT + \
                         label_data[row, col, box * PARAM_NUM + 1]
                    w = label_data[row, col, box * PARAM_NUM + 2] * CELL_WIDTH
                    h = label_data[row, col, box * PARAM_NUM + 3] * CELL_HEIGHT

                    boxes.append((cx, cy, w, h, iou))

    return boxes


def drawBoxes(image, boxes, color=(255, 0, 0), line_width=2):
    '''Draw boxes on image
    '''
    font_face = cv2.FONT_HERSHEY_COMPLEX
    font_scale = 0.5
    for box in boxes:
        cx = int(box[0] + 0.5)
        cy = int(box[1] + 0.5)
        w = int(box[2] + 0.5)
        h = int(box[3] + 0.5)
        iou_str = '%.2f' % (box[4])
        cv2.putText(image, iou_str, (cx - w // 2, cy - h // 2),
                    font_face, font_scale, color)
        cv2.rectangle(image, (cx - w // 2, cy - h // 2),
                      (cx + w // 2, cy + h // 2), color, line_width)

    return image


def drawBoxes_xyxy(image, boxes_xyxy, color=(255, 0, 0), line_width=2):
    '''Draw boxes on image
    '''
    font_face = cv2.FONT_HERSHEY_COMPLEX
    font_scale = 0.5
    for box_xyxy in boxes_xyxy:
        x1 = int(box_xyxy[0] + 0.5)
        y1 = int(box_xyxy[1] + 0.5)
        x2 = int(box_xyxy[2] + 0.5)
        y2 = int(box_xyxy[3] + 0.5)
        iou_str = '%.2f' % (box_xyxy[4])
        cv2.putText(image, iou_str, (x1, y1), font_face, font_scale, color)
        cv2.rectangle(image, (x1, y1), (x2, y2), color, line_width)

    return image


def getImageRois(raw_image_size, small_image_size, overlap_ratio=0.15):
    '''将输入的任意大小的遥感图像裁剪为符合网络输入大小的图片集合，并返回该图片集合
    Parameters:
        raw_image_size: 原始遥感图像的大小 (w, h) 格式
        small_image_size: 网络指定的输入大小 (w, h) 格式
        overlap_ratio: 相邻两张图之间的重叠率
    :return: 保存每张小图在原图中 roi 的列表
    '''

    raw_w, raw_h = raw_image_size
    s_w, s_h = small_image_size

    overlap_hor = round(s_w * overlap_ratio)
    overlap_ver = round(s_h * overlap_ratio)
    step_hor = s_w - overlap_hor
    step_ver = s_h - overlap_ver

    def makeIndexes(begin_index, end_index, interval, overlap):
        idx = [index for index in range(begin_index, end_index, interval)]
        last_index = end_index - (interval + overlap)
        if (end_index % interval) < overlap:
            idx.pop()
        idx[-1] = last_index

        return idx

    y_indexes = makeIndexes(0, raw_h, step_ver, overlap_ver)
    x_indexes = makeIndexes(0, raw_w, step_hor, overlap_hor)

    small_image_rois = []
    for x in x_indexes:
        for y in y_indexes:
            # yield (x, y, s_w, s_h)
            small_image_rois.append((x, y, s_w, s_h))

    return small_image_rois


def cuttingImage(raw_image, small_image_w_h, overlap_percentage):
    raw_height, raw_width = raw_image.shape[:2]
    small_width, small_height = small_image_w_h
    bottom = raw_height - small_height
    right = raw_width - small_width
    if bottom < 0 or right < 0:
        raw_image = cv2.copyMakeBorder(raw_image, 0, abs(bottom) * (bottom < 0),
                                       0, abs(right) * (right < 0), cv2.BORDER_CONSTANT, None, (0, 0, 0))

    raw_height, raw_width = raw_image.shape[:2]
    print((raw_width, raw_height))
    small_image_rois = getImageRois((raw_width, raw_height),
                                    small_image_w_h, overlap_percentage)

    small_images = {}
    for x, y, w, h in small_image_rois:
        small_images[(x, y)] = raw_image[y:y + h, x:x + w]

    return raw_image, small_images


def combineImages(small_images, raw_image_w_h):
    assert isinstance(small_images, dict)

    raw_width, raw_height = raw_image_w_h
    raw_image = np.zeros((raw_height, raw_width, 3), dtype=np.uint8)
    h, w = small_images[(0, 0)].shape[:2]

    for (ox, oy), small_image in small_images.items():
        raw_image[oy:oy + h, ox:ox + w] = small_image

    return raw_image


def readRectFromRawLabelTxtFile(targets, label_path):
    """
    读取 'label_path' 文件中的数据，从中提取类别为 target 的目标的位置信息并返回

    :param targets: [list] 包含需要提取的类别的列表
    :param label_path: [str] label 文件的路径
    :return: [list] label 文件中 target 的位置
    """
    assert isinstance(targets, (list, tuple))
    if not path.exists(label_path):
        return []

    f = open(label_path)
    LABEL_NAME_INDEX = 8
    targets_boxes_xyxy = [[]] * len(targets)
    for line in f.readlines():
        elements = line.split()
        if len(elements) <= LABEL_NAME_INDEX:
            continue
        category_name = elements[LABEL_NAME_INDEX]
        if category_name in targets:
            category_index = targets.index(category_name)
            elements = elements[:LABEL_NAME_INDEX + 1]
            # 将 minAreaRect 转换为 boundingBox
            points = [(int(elements[i]), int(elements[i + 1])) for i in range(0, len(elements)-1, 2)]
            x, y, w, h = cv2.boundingRect(np.asarray(points))
            targets_boxes_xyxy[category_index].append([x, y, x + w, y + h, 1])

    f.close()
    return targets_boxes_xyxy if len(targets_boxes_xyxy[0]) > 0 else []


if __name__ == "__main__":
    import cv2

    CELL_NX = 13
    CELL_NY = 13
    BOX_NUM = 4
    PARAM_NUM = 5
    IMAGE_WIDTH = 416
    IMAGE_HEIGHT = 416

    dir_path = "../../datas/train/yolt/"
    begin_index = 0
    num = -1
    image_paths = getFilePathsInDir(dir_path, ".png", begin_index, num)
    txt_paths = getFilePathsInDir(dir_path, ".txt", begin_index, num)

    for (image_path, txt_path) in zip(image_paths, txt_paths):
        print(image_path)
        print(txt_path)
        img = cv2.imread(image_path)
        label = getLabelTxtData(
            txt_path, CELL_NX, CELL_NY, BOX_NUM * PARAM_NUM)
        boxes = extractBoundingBoxesFromLabel(
            label, IMAGE_WIDTH, IMAGE_HEIGHT, BOX_NUM)
        drawBoxes(img, boxes)
        cv2.imshow("Image", img)
        # img = np.ndarray(img, np.uint8)
        # print(img.shape)
        if cv2.waitKey(0) & 0xFF == ord('q'):
            break